package com.xiashitech.agent.instrumentation.biz.userinfo.httpsession.inst;

import com.xiashitech.agent.cache.status.CacheStatus;
import com.xiashitech.agent.config.StartConfig;
import com.xiashitech.agent.constants.SpecialConfigString;
import com.xiashitech.agent.instrumentation.biz.userinfo.httpsession.SessionData;
import com.xiashitech.agent.instrumentation.utils.BaggageUtil;
import com.xiashitech.agent.utils.DurationPrinter;
import com.xiashitech.interfaces.agent.dto.AgentConfigDTO;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Scope;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

import java.util.Map;

import static com.xiashitech.agent.instrumentation.biz.userinfo.httpsession.XSHttpSessionSingletons.instrumenter;
import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext;
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
import static net.bytebuddy.matcher.ElementMatchers.*;

public class XSHttpSessionInstrumentation implements TypeInstrumentation {

    @Override
    public ElementMatcher.Junction<ClassLoader> classLoaderOptimization() {
        return hasClassesNamed("javax.servlet.http.HttpSession");
    }

    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
        return implementsInterface(named("javax.servlet.http.HttpSession"));
    }

    @Override
    public void transform(TypeTransformer transformer) {
        transformer.applyAdviceToMethod(
                named("getAttribute")
                        .and(takesArgument(0, named("java.lang.String")))
                        .and(isPublic()),
                this.getClass().getName() + "$HttpSessionAdvice");
    }

    @SuppressWarnings("unused")
    public static class HttpSessionAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(@Advice.Argument(0) String key,
                                   @Advice.Local("otelContext") Context context,
                                   @Advice.Local("otelScope") Scope scope,
                                   @Advice.Local("sessionData") SessionData sessionData) {
            DurationPrinter durationPrinter = new DurationPrinter("XSHttpSessionInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            // return 会导致 trace 数据不连续。
            //if (XSUtil.isNotEmpty(Java8Bridge.currentForBaggage().getEntryValue("xs.userId"))){
                //return;
            //}
            if (key == null || "".equals(key)){
                return;
            }
            try {
                // 过滤掉不比关心的读取，userSessionKey 是我们关心的 key
                String userSessionKey = null;
                Map<String, String> userSessionKeyMap = AgentConfigDTO.getConfig().getUserSessionKey();
                if (userSessionKeyMap != null){
                    userSessionKey = userSessionKeyMap.get(StartConfig.getParamValueFromKey(StartConfig.systemId));
                }
                if(!SpecialConfigString.ALL_SESSION.getConfiString().equalsIgnoreCase(userSessionKey) && !key.equalsIgnoreCase(userSessionKey)) {
                    return;
                }
                sessionData = SessionData.builder().key(key).build();
                Context parentContext = currentContext();
                if (!instrumenter().shouldStart(parentContext,sessionData)) {
                    return;
                }
                parentContext = BaggageUtil.syncBaggage().build().storeInContext(parentContext);
                context = instrumenter().start(parentContext, sessionData);
                context = BaggageUtil.syncBaggage().build().storeInContext(context);
                scope = context.makeCurrent();

            }catch (Throwable e){
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("HttpSessionAdvice b");
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(@Advice.Return Object value,
                                  @Advice.Argument(0) String key,
                                  @Advice.Local("otelContext") Context context,
                                  @Advice.Local("otelScope") Scope scope,
                                  @Advice.Local("sessionData") SessionData sessionData) {
            DurationPrinter durationPrinter = new DurationPrinter("XSHttpSessionInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            // return 会导致 trace 数据不连续。
            //if (XSUtil.isNotEmpty(Java8Bridge.currentForBaggage().getEntryValue("xs.userId"))){
                //return;
            //}
            if (scope == null) {
                return;
            }
            if (key == null || "".equals(key)){
                return;
            }
            try {
                // 过滤掉不比关心的读取，userSessionKey 是我们关心的 key
                String userSessionKey = null;
                Map<String, String> userSessionKeyMap = AgentConfigDTO.getConfig().getUserSessionKey();
                if (userSessionKeyMap != null){
                    userSessionKey = userSessionKeyMap.get(StartConfig.getParamValueFromKey(StartConfig.systemId));
                }
                if(!SpecialConfigString.ALL_SESSION.getConfiString().equalsIgnoreCase(userSessionKey) && !key.equalsIgnoreCase(userSessionKey)) {
                    return;
                }
                sessionData.setValue(value);
                scope.close();
                instrumenter().end(context,sessionData,sessionData,null);
            }catch (Throwable e){
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("HttpSessionAdvice e");
        }
    }

}
